import matplotlib.pyplot as plt
import numpy as np
from sklearn.linear_model import LinearRegression
font = {'family': 'Arial', 'size' : 16}
plt.rc('font', **font)

plt.rcParams['mathtext.fontset'] = 'custom'
plt.rcParams['mathtext.it'] = 'Arial:italic'
plt.rcParams['mathtext.rm'] = 'Arial'
plt.rcParams['pdf.fonttype'] = 42

# Fixing random state for reproducibility
np.random.seed(19680801)

x = [[30.47, 10.18, 20.56, 34.92, 32.49, 24.12, 32.89, 32.87],
     [27.11, 6.25, 6.44, 34.18, 7.95, 10.04, 29.50,
         27.40, 34.34, 10.02, 10.19, 32.14, 6.13, 25.84],
     [31.09, 10.32, 25.68, 10.41, 6.83, 6.06, 15.76, 22.22, 11.15, 28.26,
         16.45, 6.36, 5.92, 30.44, 24.00, 22.66, 6.66, 21.98, 9.59],
     [6.23, 20.47, 25.64, 32.45, 6.09, 23.77, 28.12, 32.26, 19.71, 29.03, 6.21, 17.35, 23.28, 10.10]]  # SBVCO

y = [[0.99, 1.37, 1.04, 0.86, 1.00, 1.08, 0.97, 0.78],
     [1.56, 2.52, 1.86, 0.72, 2.25, 2.22, 1.14,
         1.09, 0.63, 1.10, 1.98, 0.61, 2.57, 1.51],
     [0.812, 1.935, 1.524, 1.574, 2.516, 2.599, 1.42, 1.97, 1.54, 0.872,
      2.034, 2.891, 3.017, 0.835, 1.013, 1.421, 2.967, 0.892, 2.63],
     [3.64, 1.22, 1.37, 1.12, 3.61, 1.58, 0.83, 0.66, 2.43, 1.11, 3.39, 1.81, 1.68, 2.49]]  # Ea

fig, axs = plt.subplots(1, 4, constrained_layout=True, sharey=True, figsize=(12, 4))

axs[0].set_ylabel(r"$E_a (eV)$")

axs[0].text(7, 3.5, '(a)')
axs[1].text(7, 3.5, '(b)')
axs[2].text(7, 3.5, '(c)')
axs[3].text(7, 3.5, '(d)')

axs[0].annotate("(a)", xy=(0, 0.5), xytext=(-0.4, 1.3))
axs[1].annotate("(b)", xy=(0, 0.5), xytext=(-0.4, 1.3))
# for ax in axs.flat:
#     ax.set_xlabel("Step #")
#     ax.set_xlim(0, 6)
#     ax.set_xticks([0, 1, 2, 3, 4, 5, 6], labels=['0', '1', '2', '3', '2', '1', '0'])
#     # ax.set_yticks([0.5, 1.5, 2.5, 3.5], labels=['None', 'Low', 'Med', 'High'])
colors = ["r", "b", "y", "g"]
markers = ["D", "^", "s", "o"]
names = [r"$Fe$", r"$Fe_3C$", r"$Fe_5C_2$", r"$Fe_2C$"]
CFe = [0, 0.33, 0.4, 0.5]
linear = LinearRegression()
for i in range(4):
    axs[i].scatter(x[0], y[0], c=colors[0], marker=markers[0], alpha=0.1)
    axs[i].scatter(x[1], y[1], c=colors[1], marker=markers[1], alpha=0.1)
    axs[i].scatter(x[2], y[2], c=colors[2], marker=markers[2], alpha=0.1)
    axs[i].scatter(x[3], y[3], c=colors[3], marker=markers[3], alpha=0.1)
    axs[i].scatter(x[i], y[i], c=colors[i], marker=markers[i], )
    axs[i].set_xlabel(r"$SBV_{co}$")

    x_new = np.array(x[i]).reshape(-1, 1)
    y_new = np.array(y[i]).reshape(-1, 1)
    linear.fit(x_new, y_new)
    w = linear.coef_[0][0]
    b = linear.intercept_[0]
    x_line_min = min(x_new)
    x_line_max = max(x_new)
    y_line_min = w*x_line_min+b
    y_line_max = w*x_line_max+b

    axs[i].plot([x_line_min, x_line_max], [y_line_min, y_line_max], '--')
    axs[i].text(23, 3.5, names[i])
    axs[i].text(23, 3.2, r"$C/Fe=$"+"%.2f"%(CFe[i]))
    axs[i].text(23, 2.9, r"$k=$"+"%.2f"%(w))
    axs[i].text(23, 2.6, r"$b=$"+"%.2f"%(b))


# for ax in axs.flat:
#     ax.set_ylim(ax.get_ylim()[0], ax.get_ylim()[1])
#     ax.plot([3, 3], [ax.get_ylim()[0], ax.get_ylim()[1]], '--', color="Grey")
#     # ax.grid()


plt.show()